// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/exchange.h>

#include <pollux/exec/task.h>

namespace kumo::pollux::exec {
    namespace {
        std::unique_ptr<VectorSerde::Options> getVectorSerdeOptions(
            const core::QueryConfig &queryConfig,
            VectorSerde::Kind kind) {
            std::unique_ptr<VectorSerde::Options> options =
                    kind == VectorSerde::Kind::kPresto
                        ? std::make_unique<serializer::presto::PrestoVectorSerde::PrestoOptions>()
                        : std::make_unique<VectorSerde::Options>();
            options->compressionKind =
                    common::stringToCompressionKind(queryConfig.shuffleCompressionKind());
            return options;
        }
    } // namespace

    Exchange::Exchange(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::ExchangeNode> &exchangeNode,
        std::shared_ptr<ExchangeClient> exchangeClient,
        const std::string &operatorType)
        : SourceOperator(
              driverCtx,
              exchangeNode->outputType(),
              operatorId,
              exchangeNode->id(),
              operatorType),
          preferredOutputBatchBytes_{
              driverCtx->queryConfig().preferredOutputBatchBytes()
          },
          serdeKind_{exchangeNode->serdeKind()},
          serdeOptions_{
              getVectorSerdeOptions(
                  operatorCtx_->driverCtx()->queryConfig(),
                  serdeKind_)
          },
          processSplits_{operatorCtx_->driverCtx()->driverId == 0},
          driverId_{driverCtx->driverId},
          exchangeClient_{std::move(exchangeClient)} {
    }

    void Exchange::addRemoteTaskIds(std::vector<std::string> &remoteTaskIds) {
        std::shuffle(std::begin(remoteTaskIds), std::end(remoteTaskIds), rng_);
        for (const std::string &taskId: remoteTaskIds) {
            exchangeClient_->addRemoteTaskId(taskId);
        }
        stats_.wlock()->numSplits += remoteTaskIds.size();
    }

    bool Exchange::getSplits(ContinueFuture *future) {
        if (!processSplits_) {
            return false;
        }
        if (noMoreSplits_) {
            return false;
        }
        std::vector<std::string> remoteTaskIds;
        for (;;) {
            exec::Split split;
            auto reason = operatorCtx_->task()->getSplitOrFuture(
                operatorCtx_->driverCtx()->splitGroupId, planNodeId(), split, *future);
            if (reason == BlockingReason::kNotBlocked) {
                if (split.hasConnectorSplit()) {
                    auto remoteSplit = std::dynamic_pointer_cast<RemoteConnectorSplit>(
                        split.connectorSplit);
                    POLLUX_CHECK_NOT_NULL(remoteSplit, "Wrong type of split");
                    remoteTaskIds.push_back(remoteSplit->taskId);
                } else {
                    addRemoteTaskIds(remoteTaskIds);
                    exchangeClient_->noMoreRemoteTasks();
                    noMoreSplits_ = true;
                    if (atEnd_) {
                        operatorCtx_->task()->multipleSplitsFinished(
                            false, stats_.rlock()->numSplits, 0);
                        recordExchangeClientStats();
                    }
                    return false;
                }
            } else {
                addRemoteTaskIds(remoteTaskIds);
                return true;
            }
        }
    }

    BlockingReason Exchange::isBlocked(ContinueFuture *future) {
        if (!currentPages_.empty() || atEnd_) {
            return BlockingReason::kNotBlocked;
        }

        // Start fetching data right away. Do not wait for all splits to be available.

        if (!splitFuture_.valid()) {
            getSplits(&splitFuture_);
        }

        ContinueFuture dataFuture;
        currentPages_ = exchangeClient_->next(
            driverId_, preferredOutputBatchBytes_, &atEnd_, &dataFuture);
        if (!currentPages_.empty() || atEnd_) {
            if (atEnd_ && noMoreSplits_) {
                const auto numSplits = stats_.rlock()->numSplits;
                operatorCtx_->task()->multipleSplitsFinished(false, numSplits, 0);
            }
            recordExchangeClientStats();
            return BlockingReason::kNotBlocked;
        }

        // We have a dataFuture and we may also have a splitFuture_.

        if (splitFuture_.valid()) {
            // Block until data becomes available or more splits arrive.
            std::vector<ContinueFuture> futures;
            futures.push_back(std::move(splitFuture_));
            futures.push_back(std::move(dataFuture));
            *future = melon::collectAny(futures).unit();
            return BlockingReason::kWaitForSplit;
        }

        // Block until data becomes available.
        *future = std::move(dataFuture);
        return BlockingReason::kWaitForProducer;
    }

    bool Exchange::isFinished() {
        return atEnd_ && currentPages_.empty();
    }

    RowVectorPtr Exchange::getOutput() {
        if (currentPages_.empty()) {
            return nullptr;
        }

        uint64_t rawInputBytes{0};
        vector_size_t resultOffset = 0;
        if (getSerde()->supportsAppendInDeserialize()) {
            for (const auto &page: currentPages_) {
                rawInputBytes += page->size();

                auto inputStream = page->prepareStreamForDeserialize();

                while (!inputStream->atEnd()) {
                    getSerde()->deserialize(
                        inputStream.get(),
                        pool(),
                        outputType_,
                        &result_,
                        resultOffset,
                        serdeOptions_.get());
                    resultOffset = result_->size();
                }
            }
        } else {
            POLLUX_CHECK(
                getSerde()->kind() == VectorSerde::Kind::kCompactRow ||
                getSerde()->kind() == VectorSerde::Kind::kUnsafeRow);

            std::unique_ptr<melon::IOBuf> mergedBufs;
            for (const auto &page: currentPages_) {
                rawInputBytes += page->size();
                if (mergedBufs == nullptr) {
                    mergedBufs = page->getIOBuf()->clone();
                } else {
                    mergedBufs->appendToChain(page->getIOBuf()->clone());
                }
            }
            POLLUX_CHECK_NOT_NULL(mergedBufs);
            auto mergedPages = std::make_unique<SerializedPage>(std::move(mergedBufs));
            auto inputStream = mergedPages->prepareStreamForDeserialize();
            getSerde()->deserialize(
                inputStream.get(),
                pool(),
                outputType_,
                &result_,
                resultOffset,
                serdeOptions_.get());
            // We expect the row-wise deserialization to consume all the input into one
            // output vector.
            POLLUX_CHECK(inputStream->atEnd());
        }
        currentPages_.clear(); {
            auto lockedStats = stats_.wlock();
            lockedStats->rawInputBytes += rawInputBytes;
            lockedStats->rawInputPositions += result_->size();
            lockedStats->addInputVector(result_->estimate_flat_size(), result_->size());
        }

        return result_;
    }

    void Exchange::close() {
        SourceOperator::close();
        currentPages_.clear();
        result_ = nullptr;
        if (exchangeClient_) {
            recordExchangeClientStats();
            exchangeClient_->close();
        }
        exchangeClient_ = nullptr; {
            auto lockedStats = stats_.wlock();
            lockedStats->addRuntimeStat(
                Operator::kShuffleSerdeKind,
                RuntimeCounter(static_cast<int64_t>(serdeKind_)));
            lockedStats->addRuntimeStat(
                Operator::kShuffleCompressionKind,
                RuntimeCounter(static_cast<int64_t>(serdeOptions_->compressionKind)));
        }
    }

    void Exchange::recordExchangeClientStats() {
        if (!processSplits_) {
            return;
        }

        auto lockedStats = stats_.wlock();
        const auto exchangeClientStats = exchangeClient_->stats();
        for (const auto &[name, value]: exchangeClientStats) {
            lockedStats->runtimeStats.erase(name);
            lockedStats->runtimeStats.insert({name, value});
        }

        auto backgroundCpuTimeMs =
                exchangeClientStats.find(ExchangeClient::kBackgroundCpuTimeMs);
        if (backgroundCpuTimeMs != exchangeClientStats.end()) {
            const CpuWallTiming backgroundTiming{
                static_cast<uint64_t>(backgroundCpuTimeMs->second.count),
                0,
                static_cast<uint64_t>(backgroundCpuTimeMs->second.sum) *
                Timestamp::kNanosecondsInMillisecond
            };
            lockedStats->backgroundTiming.clear();
            lockedStats->backgroundTiming.add(backgroundTiming);
        }
    }

    VectorSerde *Exchange::getSerde() {
        return getNamedVectorSerde(serdeKind_);
    }
} // namespace kumo::pollux::exec
